# -*- coding:utf8 -*-
import yaml
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment
from openpyxl.styles import colors


# pip install pyyaml pandas openpyxl
# with open('./run_telechat_12b_910b.yaml', 'r', encoding='utf-8') as f:
#     result = yaml.load(f.read(), Loader=yaml.FullLoader)
#
# for k, v in result.items():
#     print(f"{k}:{v}")
# print(f"gradient_accumulation_steps:{result.get('runner_config').get('gradient_accumulation_steps')}")


class YamlToExcel(object):

    def __init__(self, input_yaml_path):
        # 实例化workbook
        self.work_book = Workbook()
        # 激活 worksheet
        self.work_sheet = self.work_book.active
        self.npu_config = self.parse(input_yaml_path)

    # 设置sheet表名
    def set_worksheet_title(self, manufacturer):
        self.work_sheet.title = manufacturer

    @classmethod
    def parse(cls, yaml_path):
        with open(yaml_path) as f:
            data = f.read()
        return yaml.load(data, Loader=yaml.Loader)

    def write_excel(self):
        # 设置列宽
        self.work_sheet.column_dimensions['A'].width = 32.0
        self.work_sheet.column_dimensions['B'].width = 33.0
        self.work_sheet.column_dimensions['C'].width = 17.0
        align = Alignment(horizontal='center', vertical='center', wrap_text=True)
        font_color = Font(color=colors.BLACK, size=12, bold=True)
        self.set_worksheet_title('参数比较')
        row_number = 0
        for key, value in self.npu_config.items():
            print(f"{key}:{value}")
            row_number += 1
            self.work_sheet.cell(row=row_number, column=1, value=f"{key}").font = font_color
            if not isinstance(value, dict):
                self.work_sheet.cell(row=row_number, column=2, value=f"{str(value)}")
            # 第2层
            if isinstance(value, dict):
                for k, v in value.items():
                    row_number += 1
                    self.work_sheet.cell(row=row_number, column=1 + 1, value=f"{k}").font = font_color
                    if not isinstance(v, dict):
                        self.work_sheet.cell(row=row_number, column=2 + 1, value=f"{str(v)}")
                    # 第3层
                    if isinstance(v, dict):
                        for k1, v1 in v.items():
                            row_number += 1
                            self.work_sheet.cell(row=row_number, column=1 + 2, value=f"{k1}").font = font_color
                            self.work_sheet.cell(row=row_number, column=2 + 2, value=f"{v1}")

    def save_excel(self, output_path):
        self.write_excel()
        self.work_book.save(output_path)


if __name__ == '__main__':
    c = YamlToExcel("E:/电信AI_telechat/src/mindformers/research/telechat/run_telechat_12b_910b.yaml")
    c.save_excel("gpu_npu.xlsx")
    dx = {"1": "s"}
